#include "hypercube_cryto_agent.hpp"
#include "intel_itt_wrapper.hpp"
#include "math_util.hpp"
#include <map>

static std::size_t get_hash(const params<helib::BGV> &p, long m,
                            const cube_signature &cube) {
  std::hash<std::string> hash;
  std::string tag;
  if (cube.is_linear()) {
    tag = std::to_string(m) + " " + std::to_string(cube.size(0)) + " " +
          std::to_string(p.hwt) + " " + std::to_string(p.t_bits) + " " +
          std::to_string(p.ksm_ncol) + " " +
          std::to_string(p.modulus_chain_bits);
  } else {
    tag = std::to_string(m) + " " + std::to_string(cube.size(1)) + " " +
          std::to_string(cube.size(1)) + " " + std::to_string(p.hwt) + " " +
          std::to_string(p.t_bits) + " " + std::to_string(p.ksm_ncol) + " " +
          std::to_string(p.modulus_chain_bits);
  }
  return hash(tag);
}

// expected row && col -> hypercube row && col
// 1. Assure that row = row_expected
// 2. Find the minimum `col` s.t. (row, col) are "good dimension"
// m: Cyclotomic polynomial - defines phi(m)
static void find_hypercube_2d(unsigned long mlb, unsigned long mub,
                              long nrow_expected, long ncol_expected,
                              long &hypercube_nrow, long &hypercube_ncol,
                              long &m, bool minimal) {
  std::unordered_map<long, long> multiple_to_mini_col;
  long row = nrow_expected;
  long col = ncol_expected;
  do {
    int multiple = 0;
    do {
      col++;
    } while (NTL::GCD(row, col) != 1);
    do {
      do {
        multiple++;
      } while (NTL::GCD(multiple, row) != 1 || NTL::GCD(multiple, col) != 1);
      m = multiple * row * col + 1;
    } while ((!NTL::ProbPrime(m) || m < mlb) && m < mub);
    if (row * col > mub) {
      if (multiple_to_mini_col.empty()) {
        ERROR_PRINTF("Bounds on m are too tight. No such m found.\n");
      } else {
        break;
      }
    }
    if (NTL::ProbPrime(m) && m >= mlb && m <= mub) {
      auto res = multiple_to_mini_col.find(multiple);
      if (res != multiple_to_mini_col.end()) {
        multiple_to_mini_col[multiple] = std::min(res->second, col);
      } else {
        multiple_to_mini_col[multiple] = col;
      }
      if (minimal) {
        continue;
      } else {
        break;
      }
    }
  } while (true);

  // Choose optimal m
  if (minimal) {
    long minimal_m = mub + 1;
    for (auto pair : multiple_to_mini_col) {
      const long cur_k = pair.first;
      const long cur_col = pair.second;
      const long cur_m = cur_k * row * cur_col + 1;
      if (minimal_m > cur_m) {
        minimal_m = cur_m;
        col = cur_col;
      }
    }
    ASSERT_PRINTF(minimal_m != mub + 1, "No such m found.\n");
    m = minimal_m;
  }

  hypercube_nrow = row;
  hypercube_ncol = col;
}

static void find_hypercube_1d(unsigned long mlb, unsigned long mub,
                              long nslot_expected, long &hypercube_nslot,
                              long &m, bool minimal) {
  std::map<long, long> multiple_to_min_nslot;
  long nslot = nslot_expected - 1;
  do {
    int multiple = 0;
    nslot++;
    do {
      do {
        multiple++;
      } while (NTL::GCD(multiple, nslot) != 1);
      m = multiple * nslot + 1;
    } while ((!NTL::ProbPrime(m) || m < mlb) && m < mub);
    if (nslot > mub) {
      if (multiple_to_min_nslot.empty()) {
        ERROR_PRINTF("Bounds on m are too tight. No such m found.\n");
      } else {
        break;
      }
    }
    if (NTL::ProbPrime(m) && m >= mlb && m <= mub) {
      auto res = multiple_to_min_nslot.find(multiple);
      if (res != multiple_to_min_nslot.end()) {
        multiple_to_min_nslot[multiple] = std::min(res->second, nslot);
      } else {
        multiple_to_min_nslot[multiple] = nslot;
      }
      if (minimal) {
        continue;
      } else {
        break;
      }
    }
  } while (true);

  // Choose optimal m
  if (minimal) {
    auto first_it = multiple_to_min_nslot.begin();
    m = first_it->first * first_it->second + 1;
  }

  hypercube_nslot = nslot;
}

std::shared_ptr<helib::Context>
hypercube_hmmcs::build_context(const params<helib::BGV> &params, long m) {
  // Hamming weight of the secret key
  const unsigned long hwt = params.hwt;
  // Number of modulus_chain_bits of the modulus chain
  const unsigned long modulus_chain_bits = params.modulus_chain_bits;
  // Number of columns of Key-Switching matrix (default = 2 or 3)
  const unsigned long ksm_ncol = params.ksm_ncol;
  // Plaintext prime modulus
  long t;
  // Bit length low bound of plaintext prime modulus
  const unsigned long t_bits = params.t_bits;
  // Hensel lifting (default = 1)
  const unsigned long r = 1;

  // Find plaintext modulus s.t. ord(t) = m / (h_nrow * h_ncol)
  find_params_time_ -= omp_get_wtime(); //! Find t & generators - Start
  std::vector<long> candidates = hypercube_.dims();
  find_ti_auto(t, t_bits, m, candidates);

  // Find generators
  std::vector<long> gens;
  std::vector<long> ords;
  find_generators(gens, ords, m, t, candidates);
  find_params_time_ += omp_get_wtime(); //! Find t & generators - Stop

  build_context_time_ -= omp_get_wtime(); //! build context - Start
  auto context =
      std::shared_ptr<helib::Context>(helib::ContextBuilder<helib::BGV>()
                                          .m(m)
                                          .p(t)
                                          .r(r)
                                          .bits(modulus_chain_bits)
                                          .skHwt(hwt)
                                          .gens(gens)
                                          .ords(ords)
                                          .c(ksm_ncol)
                                          .buildPtr());
  build_context_time_ += omp_get_wtime(); //! build context - Stop
  return context;
}

void hypercube_hmmcs::report_context(std::ostream &os) const {
  if (hypercube_.is_linear()) {
    os << "====================== Hypercube ======================\n";
    os << "* Hypercube type : Linear" << std::endl;
    os << "* # of slots     : " << hypercube_.size(0) << std::endl;
    os << "* Multiple       : " << context_->getOrdP() << std::endl;
    os << "=======================================================\n";
  } else {
    os << "====================== Hypercube ======================\n";
    os << "* Hypercube type : 2D" << std::endl;
    os << "* Rows           : " << hypercube_.size(0) << std::endl;
    os << "* Cols           : " << hypercube_.size(1) << std::endl;
    os << "* Multiple       : " << context_->getOrdP() << std::endl;
    os << "=======================================================\n";
  }
  hmm_crypto_agent::report_context();
}

void hypercube_hmmcs::report_time() const {
  // clang-format off
  std::cout << "* Setup Time                 : " << setup_time_ << " s\n";
  std::cout << "  * Find params              : " << find_params_time_ * 1000.0 << " ms\n";
  std::cout << "  * Build context            : " << build_context_time_ << " s\n";
  // clang-format on
}

hypercube_hmmcs::hypercube_hmmcs(const params<helib::BGV> &params,
                                 const cube_signature &cube, bool cache_free,
                                 bool minimal_m, bool verbose) {
  // ITT API
  INTEL_ITT_DOMAIN_CREATE(hyper_hmme_domain, "src.hypercube.hypercube_hmmcs");
  INTEL_ITT_STRING_HANDLE_CREATE(setup_handle, "build context");

  const unsigned long mlb = params.m_lower_bound;
  const unsigned long mub = params.m_upper_bound;

  setup_time_ -= omp_get_wtime(); //! Setup - Start
  long m;
  find_params_time_ -= omp_get_wtime(); //! Find m - Start
  if (cube.is_linear()) {
    long nslot;
    find_hypercube_1d(mlb, mub, cube.size(0), nslot, m, minimal_m);
    // Init hypercube
    hypercube_.init({nslot});
  } else {
    long h_nrow, h_ncol;
    find_hypercube_2d(mlb, mub, cube.size(0), cube.size(1), h_nrow, h_ncol, m,
                      minimal_m);
    // Init hypercube
    hypercube_.init({h_nrow, h_ncol});
  }
  find_params_time_ += omp_get_wtime(); //! Find m - Stop
  if (cache_free) {
    COND_INFO_PRINTF(verbose, "Not using context cache\n");
    INTEL_ITT_TASK_BEGIN(hyper_hmme_domain, setup_handle); //! ITT - Begin
    context_ = build_context(params, m);
    INTEL_ITT_TASK_END(hyper_hmme_domain); //! ITT - End
  } else {
    const auto uuid = get_hash(params, m, hypercube_);
    const std::string name = "context_" + std::to_string(uuid) + ".cache";
    std::ifstream in(name);
    if (!in.is_open()) {
      // No context: rebuild
      COND_INFO_PRINTF(verbose, "No context cache for this crypto(%s)\n",
                       name.c_str());
      INTEL_ITT_TASK_BEGIN(hyper_hmme_domain, setup_handle); //! ITT - Begin
      context_ = build_context(params, m);
      INTEL_ITT_TASK_END(hyper_hmme_domain); //! ITT - End

      // Save context
      std::ofstream of(name);
      ASSERT_PRINTF(of.is_open(), "Can not open %s\n", name.c_str());
      context_->writeTo(of);
      COND_INFO_PRINTF(verbose, "Save current context in %s\n", name.c_str());
    } else {
      // Load serailized context
      COND_INFO_PRINTF(verbose, "Load serialized context in %s\n",
                       name.c_str());
      INTEL_ITT_TASK_BEGIN(hyper_hmme_domain, setup_handle); //! ITT - Begin
      build_context_time_ -= omp_get_wtime(); //! Build context - Start
      context_ =
          std::shared_ptr<helib::Context>(helib::Context::readPtrFrom(in));
      build_context_time_ += omp_get_wtime(); //! Build context - Stop
      INTEL_ITT_TASK_END(hyper_hmme_domain);  //! ITT - End
    }
  }
  setup_time_ += omp_get_wtime(); //! Setup - Stop
}

hypercube_hmmcc::hypercube_hmmcc(const params<helib::BGV> &params,
                                 const cube_signature &cube, bool cache_free,
                                 bool minimal_m)
    : hypercube_hmmcs(params, cube, cache_free, minimal_m, true),
      hmm_crypto_client(*context_) {
  //* Construct `hypercube_hmmcs` instance to build the HElib context
  //* Using `context` to construct `hmm_crypto_client` instance

  setup_time_ += key_gen_time_; //! Part of Setup
}

hypercube_hmmcc::hypercube_hmmcc(const hypercube_hmmcs &hmmcs)
    : hypercube_hmmcs(hmmcs.context(), hmmcs.build_context_time()),
      hmm_crypto_client(*context_) {}

void hypercube_hmmcc::report_time() const {
  hypercube_hmmcs::report_time();
  std::cout << "  * KeyGen Time              : " << key_gen_time_ << " s\n";
}

//* Encryption & Decryption
static inline void pack_cube(helib::Ptxt<helib::BGV> &ptxt,
                             const NTL::mat_ZZ &mat,
                             const helib::EncryptedArray &ea) {
  const long row = mat.NumRows();
  const long col = mat.NumCols();
  if (ea.dimension() == 2) {
    long h_row = ea.sizeOfDimension(0);
    long h_col = ea.sizeOfDimension(1);
    for (auto i = 0; i < h_row; i++) {
      for (auto j = 0; j < h_col; j++) {
        int idx = i * h_col + j;
        if (j >= col || i >= row) {
          // 1. Exceeding the actual # of rows, or
          // 2. Exceeding the actual # of columns
          ptxt[idx] = 0;
        } else {
          ptxt[idx] = NTL::conv<long>(mat[i][j]);
        }
      }
    }
  } else {
    for (auto i = 0; i < row; i++) {
      for (auto j = 0; j < col; j++) {
        int idx = i * col + j;
        ptxt[idx] = NTL::conv<long>(mat[i][j]);
      }
    }
  }
}

static inline NTL::mat_ZZ unpack_cube(const helib::Ptxt<helib::BGV> &ptxt,
                                      const helib::EncryptedArray &ea) {
  NTL::mat_ZZ A;
  if (ea.dimension() == 2) {
    long h_row = ea.sizeOfDimension(0);
    long h_col = ea.sizeOfDimension(1);
    A.SetDims(h_row, h_col);
    for (auto i = 0; i < h_row; i++) {
      for (auto j = 0; j < h_col; j++) {
        int idx = i * h_col + j;
        A[i][j] = NTL::coeff(ptxt[idx].getData(), 0);
      }
    }
  } else if (ea.dimension() == 1) {
    long nslot = ea.sizeOfDimension(0);
    A.SetDims(1, nslot);
    for (auto i = 0; i < nslot; i++) {
      A[0][i] = NTL::coeff(ptxt[i].getData(), 0);
    }
  }
  return std::move(A);
}

void hypercube_hmmcc::encrypt(helib::Ctxt &ctxt, const NTL::mat_ZZ &mat,
                              ENC type) {
  // sanity check
  ASSERT_PRINTF(ctxt.getPubKey() == public_key(),
                "Encrypt: public key mismatch\n");

  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(*context_);

  // Packing matrix
  pack_cube(ptxt, mat, context_->getEA());

  hmm_crypto_client::encrypt(ctxt, ptxt, type);
}
helib::Ctxt hypercube_hmmcc::encrypt(const NTL::mat_ZZ &mat, ENC type) {
  auto ctxt = get_empty_ctxt();
  encrypt(ctxt, mat, type);
  return std::move(ctxt);
}

helib::Ctxt hypercube_hmmcc::encrypt(const std::vector<long> &vec, ENC type) {
  auto ctxt = this->get_empty_ctxt();
  helib::PtxtArray ptxt(*context_);
  ptxt.load(vec);
  ptxt.encrypt(ctxt);
  return std::move(ctxt);
}

helib::Ctxt hypercube_hmmcc::encrypt(const std::vector<double> &vec, ENC type) {
  auto ctxt = this->get_empty_ctxt();
  helib::PtxtArray ptxt(*context_);
  ptxt.load(vec);
  ptxt.encrypt(ctxt);
  return std::move(ctxt);
}

void hypercube_hmmcc::decrypt(NTL::mat_ZZ &mat, const helib::Ctxt &cmat) {
  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(*context_);
  // Decrypt
  hmm_crypto_client::decrypt(ptxt, cmat);
  mat = std::move(unpack_cube(ptxt, context_->getEA()));
}

void hypercube_hmmcc::decrypt(std::vector<double> &vec,
                              const helib::Ctxt &cmat) {
  // Create a vector of long with nslots elements
  helib::PtxtArray ptxt(*context_);
  // Decrypt
  ptxt.decrypt(cmat, *(this->sk_));
  ptxt.store(vec);
}

void hypercube_hmmcc::decrypt(std::vector<long> &vec, const helib::Ctxt &cmat) {
  // Create a vector of long with nslots elements
  helib::PtxtArray ptxt(*context_);
  // Decrypt
  ptxt.decrypt(cmat, *(this->sk_));
  ptxt.store(vec);
}